import os

import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, StorageContext
from llama_index.core.node_parser import TokenTextSplitter, SentenceSplitter
from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from llama_index.vector_stores.chroma import ChromaVectorStore

# 设置默认的大语言模型为 DashScope 的 Qwen Max 模型，并从环境变量中读取 API Key
Settings.llm = DashScope(model_name=DashScopeGenerationModels.QWEN_MAX, api_key=os.getenv("DASHSCOPE_API_KEY"))

# 替换 LlamaIndex 默认的嵌入模型为 DashScope 提供的文本嵌入模型
Settings.embed_model = DashScopeEmbedding(
    model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V3,
)

# 初始化 Chroma 客户端并配置持久化存储路径
from chromadb.config import Settings

client = chromadb.Client(
    Settings(
        persist_directory="./chroma_data",  # 数据持久化保存路径
        is_persistent=True  # 启用数据持久化功能
    )
)

collection_name = "deepseek_docs"

# 获取或创建一个名为 deepseek_docs 的集合，用于存储向量数据
# 如果该集合不存在则自动创建，并附加元信息描述
collection = client.get_or_create_collection(
    name=collection_name,
    metadata={"description": "DeepSeek V3 相关文档集合"}
)

# 配置 Chroma 向量数据库与 LlamaIndex 存储上下文之间的连接桥梁
vector_store = ChromaVectorStore(chroma_collection=collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 加载指定目录下的 PDF 文件作为输入文档
documents = SimpleDirectoryReader(
    "./data",
    required_exts=[".pdf"],
).load_data()

# 使用 TokenTextSplitter 对加载的文档进行切分，便于后续构建节点
node_parser = TokenTextSplitter(chunk_size=512, chunk_overlap=200)

# 将原始文档转换为可索引的节点结构
nodes = node_parser.get_nodes_from_documents(documents)

# 构建基于向量存储的索引结构，将节点存入 Chroma 中
index = VectorStoreIndex(nodes, storage_context=storage_context)

# 创建检索器对象，用于执行语义相似度查询，返回最相关的 top-k 个结果
vector_retriever = index.as_retriever(similarity_top_k=1)

# 执行一次具体的检索任务，查找关于 "deepseek v3 数学能力怎么样" 的相关内容
results = vector_retriever.retrieve("deepseek v3数学能力怎么样")

# 输出匹配到的第一个相关结果
print(results[0])
